# 这段代码定义了一个名为RKNet的神经网络模型，并且包含了一个USAM模块作为其特征提取部分。下面是代码各部分的解析：
# 1、USAM 类（模块）
# USAM 代表“Universal Spatial Attention Module”（通用空间注意力模块）。这个模块的主要作用是通过空间注意力机制增强输入特征图（feature map）。其工作流程如下：
# 初始化：
# kernel_size 和 padding：指定卷积核大小和卷积操作的填充大小。
# weight：一个常量的卷积核，用于对输入特征图的求和操作。
# weight2：一个常量的卷积核，用于计算卷积输出。
# polish：如果为True，则会将输出特征图的边界区域的注意力值置为零。
# bn：批量归一化层，用于对特征图进行归一化。
# relu：ReLU激活函数，用于非线性变换。
# 前向传播：
# 首先对输入x的每个通道求和，并得到一个单通道特征图fmap。
# 通过F.conv2d使用weight进行卷积，得到x1，然后使用weight2进行卷积得到x2。
# 两者之差（att = x2 - x1）作为空间注意力特征，进行批量归一化和ReLU激活。
# 如果polish为True，会对特征图的边界区域（四个边缘）进行“抛光”，即将这些区域的值置为零。
# 最后，将输入x和注意力特征图的加权输出相加。
#
# 2、RKNet 类
# RKNet 是一个结合了ResNet50预训练模型和USAM模块的神经网络模型。它主要用于处理图像并提取特征。
# 初始化：
# model_ft：加载ResNet50预训练模型，并进行必要的修改。如果stride为1，则会调整ResNet中的layer4的步幅（stride），保持图像尺寸。
# pool：指定池化方式。可以选择：
# 'avg'：使用自适应平均池化。
# 'max'：使用自适应最大池化。
# 'avg+max'：同时使用自适应平均池化和最大池化。
# 'gem'：使用GeM（Generalized Mean Pooling）池化。
# 根据传入的init_model参数，可能会用一个初始化的模型来替换当前的model_ft。
# 添加两个USAM模块：usam_1和usam_2，用于在网络的不同层进行空间注意力计算。
# 前向传播：
# 输入图像x经过ResNet的第一层卷积、批归一化、ReLU激活后，传递到usam_1进行空间注意力处理。
# 然后，经过ResNet的maxpool层和layer1，再经过usam_2处理，接着传递到layer2、layer3和layer4。
# 最终的输出是通过ResNet的深层网络得到的特征图。
#
# 总结
# USAM模块：通过空间注意力机制，增强了特征图的空间表达，尤其是对边界区域的抛光操作。
# RKNet模型：结合了ResNet50的强大特征提取能力和USAM的空间注意力机制，适用于图像分类、检测或其他计算机视觉任务。

from torch import nn
import torch
import torch.nn.functional as F
from torchvision import models

class USAM(nn.Module):
    def __init__(self, kernel_size=3, padding=1, polish=False):
        super(USAM, self).__init__()

        kernel = torch.ones((kernel_size, kernel_size))
        kernel = kernel.unsqueeze(0).unsqueeze(0)
        self.weight = nn.Parameter(data=kernel, requires_grad=False)
        

        kernel2 = torch.ones((1, 1)) * (kernel_size * kernel_size)
        kernel2 = kernel2.unsqueeze(0).unsqueeze(0)
        self.weight2 = nn.Parameter(data=kernel2, requires_grad=False)

        self.polish = polish
        self.pad = padding
        self.relu = nn.ReLU()
        self.bn = nn.BatchNorm2d(1)

    def __call__(self, x):
        fmap = x.sum(1, keepdim=True)      
        x1 = F.conv2d(fmap, self.weight, padding=self.pad)
        x2 = F.conv2d(fmap, self.weight2, padding=0) 
        
        att = x2 - x1
        att = self.bn(att)
        att = self.relu(att)

        if self.polish:
            att[:, :, :, 0] = 0
            att[:, :, :, -1] = 0
            att[:, :, 0, :] = 0
            att[:, :, -1, :] = 0

        output = x + att * x

        return output



class RKNet(nn.Module):
    def __init__(self, stride=2, init_model=None, pool='avg'):
        super(RKNet, self).__init__()
        model_ft = models.resnet50(pretrained=True)
        # avg pooling to global pooling
        if stride == 1:
            model_ft.layer4[0].downsample[0].stride = (1,1)
            model_ft.layer4[0].conv2.stride = (1,1)

        self.pool = pool
        if pool =='avg+max':
            model_ft.avgpool2 = nn.AdaptiveAvgPool2d((1,1))
            model_ft.maxpool2 = nn.AdaptiveMaxPool2d((1,1))
            #self.classifier = ClassBlock(4096, class_num, droprate)
        elif pool=='avg':
            model_ft.avgpool2 = nn.AdaptiveAvgPool2d((1,1))
            #self.classifier = ClassBlock(2048, class_num, droprate)
        elif pool=='max':
            model_ft.maxpool2 = nn.AdaptiveMaxPool2d((1,1))
        elif pool=='gem':
            model_ft.gem2 = GeM(dim=2048)

        self.model = model_ft

        if init_model!=None:
            self.model = init_model.model
            self.pool = init_model.pool
            #self.classifier.add_block = init_model.classifier.add_block

        self.usam_1 = USAM()
        self.usam_2 = USAM()

    def forward_features(self, x):
        x = self.model.conv1(x)
        x = self.model.bn1(x)
        x = self.model.relu(x)
        x = self.usam_1(x)
        x = self.model.maxpool(x)
        x = self.model.layer1(x)
        x = self.usam_2(x)
        x = self.model.layer2(x)
        x = self.model.layer3(x)
        x = self.model.layer4(x)
        
        return x